import torch
import torchvision.datasets as dsets
import torchvision.transforms as transforms

# 设置每次训练输入的样本数量为200
batch_size = 200
# 设置整个训练集被训练的轮数为10
num_epochs = 10
# 设置学习率为0.001
learning_rate = 0.001

# 准备数据集并进行预处理
train_dataset = dsets.FashionMNIST(root='./dataset', train=True, transform=transforms.ToTensor(), download=True)
test_dataset = dsets.FashionMNIST(root='./dataset', train=False, transform=transforms.ToTensor(), download=True)

# 构建数据加载器
# 在训练集中打乱以提高模型的泛化能力，在测试集中不打乱以保证每个样本都有相同的评估标准
train_loader = torch.utils.data.DataLoader(dataset=train_dataset, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset, batch_size=batch_size, shuffle=False)

# 定义CNN模型
class CNN(torch.nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 定义卷积层和最大池化层
        # 第一层卷积层输入通道为1，输出通道为32，卷积核大小为5*5，每次滑动步长为1个像素点
        # 在输入图像周围添加2个像素宽度的零填充，以保证卷积后输出特征图与原始图像大小相同。
        self.conv1 = torch.nn.Conv2d(in_channels=1, out_channels=32, kernel_size=5, stride=1, padding=2)
        # 第一层池化层卷积核大小为2*2
        self.pool1 = torch.nn.MaxPool2d(kernel_size=2)
        # 第二层卷积层输入通道为上个卷积层的输出通道32，输出通道为64，卷积核大小为5*5，每次滑动步长为1个像素点
        self.conv2 = torch.nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=1, padding=2)
        # 第二层池化层卷积核大小为2*2
        self.pool2 = torch.nn.MaxPool2d(kernel_size=2)
        # 定义全连接层
        # 第一层全连接层输入大小为7*7*64，输出大小为256
        self.fc1 = torch.nn.Linear(7 * 7 * 64, 256)
        # 第二层全连接层输入大小为上层输出大小256，输出大小为10，分别对应10个输出数据
        self.fc2 = torch.nn.Linear(256, 10)

    def forward(self, x):
        # 将输入张量x输入卷积层进行卷积操作，提取样本图片特征
        x = self.conv1(x)
        # 使用ReLU函数对输入张量x进行激活操作
        x = torch.nn.functional.relu(x)
        # 对输入张量x进行池化操作，提取最显著的特征，并减小特征图的大小，进一步降低计算复杂度。
        x = self.pool1(x)
        x = self.conv2(x)
        x = torch.nn.functional.relu(x)
        x = self.pool2(x)
        # 将经过两次卷积池化后的数据x展平成一维向量形式
        x = x.view(-1, 7 * 7 * 64)
        x = self.fc1(x)
        x = torch.nn.functional.relu(x)
        x = self.fc2(x)
        return x

# 创建模型实例并定义损失函数和优化器
# 创建CNN模型实例，即生成一个可以识别样本数据并输出预测结果的模型对象。
model = CNN()
# 定义损失函数为交叉熵，用于度量模型预测结果和真实标签之间的差距。
criterion = torch.nn.CrossEntropyLoss()
# 定义优化器为Adam优化器，并传入模型的所有可训练参数（model.parameters()），学习率为learning_rate。
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

# 训练模型
total_step = len(train_loader)
for epoch in range(num_epochs):
    for i, (images, labels) in enumerate(train_loader):
        # 使用model对输入值images进行预测，将预测值赋给output
        outputs = model(images)
        # 用损失函数计算预测值与标签的误差
        loss = criterion(outputs, labels)
        # 每次循环前使用optimizer.zero_grad()方法来将梯度清零，以防止梯度累积影响模型性能
        optimizer.zero_grad()
        # 反向传播将当前损失值向输出端进行传播，并计算出所有需要计算的参数的梯度值，以便进行优化和参数更新。
        loss.backward()
        # 使用optimizer.step()来更新网络参数
        optimizer.step()
        # 打印训练信息
        if (i + 1) % 100 == 0:
            print(f"Epoch [{epoch + 1}/{num_epochs}], Step [{i + 1}/{total_step}], Loss: {loss.item():.4f}")
# 评估模型性能
# 测试集不需要计算梯度
with torch.no_grad():
    correct = 0
    total = 0
    for images, labels in test_loader:
        # 将测试集数据输入训练好的模型，得到预测值，并将预测值赋给outputs
        outputs = model(images)
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
    accuracy = correct / total * 100
    print(f'测试集准确率:{accuracy}%')